from pymongo import UpdateOne
from pev2.factor.BaseFactor import BaseFactor
from pev2.data.FinanceReportCrawler import getCodeReports
from pev2.data.DataModule import DataModule

"""
实现市盈率因子的计算和保存
"""
class PeFactor(BaseFactor):

    def __init__(self):
        BaseFactor.__init__(self, name='pe')

    """
    计算指定时间段内所有股票的该因子的值，并保存到数据库中
    :param startDate:  开始时间
    :param endDate: 结束时间
    """
    def compute(self, startDate, endDate):
        dm = DataModule()

        codeReportDict = getCodeReports()

        codes = set(codeReportDict.keys())

        for code in codes:
            dailies = dm.getKData(code, autype=None, startDate=startDate, endDate=endDate)

            # 如果没有合适的数据
            if dailies.index.size == 0:
                continue

            # 业绩报告列表
            reports = codeReportDict[code]

            dailies.set_index(['date'], inplace=True)

            updateRequests = []
            for currentDate in dailies.index:
                # 用来保存最后一个公告日期小于等于当前日期的财报
                lastReport = None

                for report in reports:
                    announcedDate = report['announced_date']
                    # 如果公告日期大于当前调整日，则结束循环
                    if announcedDate > currentDate:
                        break
                    lastReport = report

                # 如果找到了正确时间范围的年报，则计算PE
                if lastReport is not None:
                    pe = dailies.loc[currentDate]['close']/lastReport['eps']
                    pe = round(pe, 3)
                    print('%s, %s, %s, eps: %5.2f, pe: %6.2f' %
                          (code, currentDate, lastReport['announced_date'], lastReport['eps'], pe),
                          flush=True)
                    updateRequests.append(
                        UpdateOne(
                            {'code': code, 'date': currentDate},
                            {'$set':{'code': code, 'date': currentDate, 'pe': pe}}, upsert=True
                        )
                    )
                if len(updateRequests) > 0:
                    saveResult = self.conn.bulk_write(updateRequests, ordered=False)
                    print('股票代码: %s, 因子: %s, 插入：%4d, 更新: %4d' %
                          (code, self.name, saveResult.upserted_count, saveResult.modified_count), flush=True)



PeFactor().compute('2014-01-01', '2014-12-31')
